package com.aiwiown.snackmq.cluster.metadata;

import com.aiwiown.snackmq.common.cluster.BrokerId;
import com.aiwiown.snackmq.common.protocol.TopicPartition;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 负责计算分区分配方案的算法类。
 */
@Slf4j
public class Rebalancer {

    /**
     * 一个简单的轮询再均衡算法。
     * 在真实的系统中，这会是一个非常复杂的算法，旨在最小化数据移动。
     * @param liveBrokers   当前所有可用的 Broker
     * @return 一个代表新的、均衡的分配方案的 Map
     */
    public Map<TopicPartition, List<BrokerId>> calculateNewAssignments(
            Map<String, TopicMetadata> topics, List<BrokerId> liveBrokers) {

        final Map<TopicPartition, List<BrokerId>> newAssignments = new HashMap<>();
        if (liveBrokers.isEmpty()) {
            return newAssignments;
        }
        topics.forEach((topicName, topicMetadata) -> {
            // 【核心修改】从元数据中获取期望的副本数
            int replicationFactor = topicMetadata.getReplicationFactor();
            // 做一个保护，确保副本数不会超过存活的 Broker 数
            int actualReplicationFactor = Math.min(replicationFactor, liveBrokers.size());

            if (actualReplicationFactor < replicationFactor) {
                log.warn("Topic '{}' 期望的副本数 ({}) 大于当前存活的 Broker 数 ({}). " +
                                "将使用 {} 作为实际副本数。",
                        topicName, replicationFactor, liveBrokers.size(), actualReplicationFactor);
            }
            if (actualReplicationFactor == 0) {
                return; // 如果没有存活节点，无法分配
            }

            for (Integer partitionId : topicMetadata.getPartitionAssignments().keySet()) {
                TopicPartition tp = new TopicPartition(topicName, partitionId);
                List<BrokerId> replicas = assignReplicasForPartition(tp, liveBrokers, actualReplicationFactor);
                newAssignments.put(tp, replicas);
            }
        });

        log.info("计算出新的分区分配方案: {}", newAssignments);
        return newAssignments;
    }

    /**
     * 为特定分区分配副本。
     * 修复：确保在故障转移场景中，副本能够正确分布在不同的 Broker 上。
     *
     * @param tp            分区信息
     * @param brokers       可用的 Broker 列表
     * @param count         需要分配的副本数量
     * @return              分配的副本列表
     */
    private List<BrokerId> assignReplicasForPartition(TopicPartition tp, List<BrokerId> brokers, int count) {
        List<BrokerId> replicas = new ArrayList<>();
        
        // 如果可用的 Broker 数量不足，使用所有可用的 Broker
        int actualCount = Math.min(count, brokers.size());
        
        // 使用主题名称和分区ID的组合来确定起始位置，确保分配的一致性
        int startIndex = Math.abs((tp.getTopic().hashCode() + tp.getPartitionId())) % brokers.size();
        
        for (int i = 0; i < actualCount; i++) {
            int brokerIndex = (startIndex + i) % brokers.size();
            BrokerId broker = brokers.get(brokerIndex);
            
            // 避免重复添加同一个 Broker
            if (!replicas.contains(broker)) {
                replicas.add(broker);
            } else if (brokers.size() > 1) {
                // 如果当前 Broker 已经存在，尝试找到下一个可用的 Broker
                for (int j = 1; j < brokers.size(); j++) {
                    int nextIndex = (brokerIndex + j) % brokers.size();
                    BrokerId nextBroker = brokers.get(nextIndex);
                    if (!replicas.contains(nextBroker)) {
                        replicas.add(nextBroker);
                        break;
                    }
                }
            }
        }
        
        return replicas;
    }
}
